import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.models import load_model

# 指定模型文件路径
model_path = r'E:\EEG\EEG-TransNet-main\checkpoints\A09T\generator_epoch_120.h5'

# 加载模型
generator = load_model(model_path)

# 打印模型概要
print(generator.summary())

# 生成随机噪声
noise_dim = 256  # 假设噪声维度为100
num_examples = 10  # 生成3个示例
random_noise = np.random.normal(size=(num_examples, noise_dim))

# 使用生成器生成 EEG 数据
generated_samples = generator.predict(random_noise)

# 可视化生成的 EEG 数据
def visualize_generated_data(generated_samples, sample_index=0):
    # 选择一个样本进行可视化
    fake_sample = generated_samples[sample_index]

    # 如果生成的样本形状为 (1, 22, 1125)，需要将其 reshape 成 (22, 1125)
    if fake_sample.shape == (1, 22, 1125):
        fake_sample = fake_sample.reshape(22, 1125)

    # 可视化伪造的 EEG 数据
    plt.figure(figsize=(15, 10))
    colors = plt.cm.viridis(np.linspace(0, 1, 22))  # 生成22种不同的颜色
    for i in range(22):
        plt.plot(fake_sample[i, :], label=f'Channel {i + 1}', color=colors[i])

    plt.title('Fake EEG Data')
    plt.xlabel('Time')
    plt.ylabel('Amplitude')
    plt.legend(loc='upper right')
    plt.show()

# 只可视化前10个生成的样本
for i in range(min(10, num_examples)):
    visualize_generated_data(generated_samples, sample_index=i)

# 保存生成的数据
save_dir = r'E:\EEG\EEG-TransNet-main\data\dataset\bci_iv_2a'
for i in range(num_examples):
    file_name = f'A09Efake_eeg_{i}.npy'
    file_path = f'{save_dir}\{file_name}'
    np.save(file_path, generated_samples[i])
    print(f'Saved {file_name} to {save_dir}')

# 可视化第一个生成的样本
visualize_generated_data(generated_samples, sample_index=0)